#%%
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, pandas as pd
%matplotlib inline

grains = ['tg', 'ng2', 'ng3', 'ng4', 'ng5', 'sf']

df_density = pd.read_csv('1_var_density_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_ars = pd.read_csv('3_var_ARS_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_vels = pd.read_csv('4_var_settling_velocities_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_VES = pd.read_csv('5_var_VES_settling_velocity_2022.csv').set_index('Unnamed: 0').rename_axis(None)

# %%
# calculate cstar and mustar

bed_angle = 3.5
aor_gb = 24.736
rof = 998.21
g = 9.81
cstar, dcstar, mustar, dmustar, ratio, dratio, Co, CDsettle, CD = {}, {}, {}, {}, {}, {}, {}, {}, {}
for grain in grains:
    cstar[grain] = ((df_VES.ws_sphere[grain])/(df_vels.vel_mean[grain]))**2 * df_shape.CSF[grain]
    dcstar[grain] = cstar[grain] * np.sqrt((2*df_VES.dws_sphere[grain]/df_VES.ws_sphere[grain])**2 + (2*df_vels.dvel_mean[grain]/df_vels.vel_mean[grain])**2 + (df_shape.dCSF[grain]/df_shape.CSF[grain])**2)
    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar[grain] = (np.tan(np.radians(df_ars.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
    dmustar[grain] = ((1 + np.tan(np.radians(df_ars.mean_angle[grain]))**2)/mugb) * np.radians(df_ars.mean_angle[grain]) * (df_ars.std_angle[grain]/df_ars.mean_angle[grain])
    ratio[grain] = cstar[grain]/mustar[grain]
    dratio[grain] = ratio[grain] * np.sqrt((dcstar[grain]/cstar[grain])**2 + (dmustar[grain]/mustar[grain])**2)
    Co[grain] = df_VES.CD_VES[grain]
    CDsettle[grain] = 4*(df_density.total_grain_ro[grain]/rof - 1)*g*df_shape.dn[grain] / (3*df_vels.vel_mean[grain]**2)
    CD[grain] = df_vels.CD[grain]

#%%
# # %%
database = {
        'mu*': mustar,
        'dmu*': dmustar,
        'C*': cstar,
        'dC*': dcstar,
        'C*/mu*': ratio,
        'd(C*/mu*)': dratio,
        'Co': Co,
        'CDsettle': CDsettle,
        'CD': CD
        }

df_all = pd.DataFrame(database)
df_all.to_csv('grain_database_2022_new.csv')
df_all
# %%
